from nets.yolo3 import yolo_body
from keras.layers import Input
from yolo import YOLO
from PIL import Image
from cv import videoToImg
from cv import mkdir
import cv2
import os.path
import os
import json





def f_realtime(video_path):
    


    def cropImg(img, box):                                    #按照尺寸裁剪图片
        cropped = img.crop((box[1], box[0], box[3], box[2]))
        return cropped

    def getCropImgPath(video, box_class):
        dir_name = 'D:/桌面/yolo3-keras/croppedImg/' + video
        if not os.path.exists(dir_name):
            os.mkdir(dir_name)
        _count = traffic_count[box_class]
        crapPath = dir_name + '/' + class_names[box_class] + str(_count) + '.jpg'
        return crapPath


    # 检测车流量
    def traffic(old_boxs, new_boxs, old_classes, new_classes, img, frame):
        for i in range(len(new_boxs)):
            current_box = new_boxs[i]                            #当前box
            current_x = (current_box[1] + current_box[3]) / 2    #求x中点
            current_y = (current_box[0] + current_box[2]) / 2    #求y中点
            current_width = current_box[3] - current_box[1]      #求宽度
            current_height = current_box[2] - current_box[0]     #求高度
            current_class = new_classes[i]                       #当前种类
            for j in range(len(old_boxs)):
                old_box = old_boxs[j]
                old_x = (old_box[1] + old_box[3]) / 2
                old_y = (old_box[0] + old_box[2]) / 2
                old_class = old_classes[j]
                # 判断是否在判定范围内，并且为相同class
                if ((current_class == old_class) and (abs(current_x - old_x) < current_width) and (abs(current_y - old_y) < current_height)):
                    # 判断是否刚超过baseline
                    if (((current_x - img_width * baseline_x) * (old_x - img_width * baseline_x) < 0)):
                        if (current_y > img_width * 0.7):
                            continue
                        traffic_count[current_class] += 1
                        cropped = cropImg(img, current_box)
                        cropPath = getCropImgPath(videoName, current_class)
                        cropped.save(cropPath)
                        print("in frame " + str(frame) + " find " + str(class_names[current_class]))
                        break
                    elif((current_y - img_height * baseline_y) * (old_y - img_height * baseline_y) < 0):
                    
                        traffic_count[current_class] += 1
                        cropped = cropImg(img, current_box)
                        cropPath = getCropImgPath(videoName, current_class)
                        cropped.save(cropPath)
                        print("in frame " + str(frame) + " find " + str(class_names[current_class]))
                        break

    # 获取所有class
    def get_class():
        classes_path = os.path.expanduser('D:/桌面/yolo3-keras/model_data/coco_classes.txt')
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

#*****************************************************************************
    old_boxs = [] #前一帧box
    old_classes = [] #前一帧class
    baseline_x = 0.3 #x基准线
    baseline_y = 0.7 #y基准线
    img_width = 0 #图片宽度
    img_height = 0 #图片高度
    traffic_x_count = [0] * 80 #横向车道每一种class计数
    traffic_y_count = [0] * 80 #纵向车道每一种class计数
    traffic_count = [0] * 80
    class_names = get_class()  #获取，每种物体的名称
    videoName = ''
    yolo = YOLO()


    
    try:
        num, img_width, img_height = videoToImg(video_path, 6)  #每6帧获取一张图片
    except:
        print('Open Error! Try again!')
    else:
        # traffic_x_count = [0] * 80 #横向车道每一种class计数
        # traffic_y_count = [0] * 80 #纵向车道每一种class计数 
        traffic_count = [0] * 80
        videoName = os.path.basename(video_path)
        inputFilePath = './output-img/' + videoName + '/'    #输入图片路径
        predictFilePath = "./predict-img/" + videoName + "/" #预测图片路径
        mkdir(predictFilePath)                               #创建预测图片路径
        object_rank = 0
        
        #创建json文件
        jsonadd = predictFilePath + videoName+'.json'        #存储所有物体的json文件的路径
        carplatej = predictFilePath + videoName+'car'+'.json'#存储识别出车牌的车辆的json文件路径
        jsonframe = predictFilePath + videoName+'jsonframe'+'.json'#存储每帧所含物体的json文件路径
        
        #创建临时json数组
        fuckjson = []                                        #存储所有物体的json文件的数组
        fuckjsoncar = []                                     #存储识别出车牌的车辆的json文件数组
        fuckframejson = []                                   #存储每帧所含物体的json文件数组
        
        #创建车牌数组
        tempcar = []
        for i in range(num):                                 #开始遍历每一帧

            imgpath = inputFilePath + str(i + 1) + '.jpg'    #输入文件路径
            image = Image.open(imgpath)                      #打开输入文件
            detected,object_rank = yolo.detect_image(image,predictFilePath,object_rank,
            fuckjson,
            fuckjsoncar,
            fuckframejson,
            tempcar,

            )                                                #对图片进行检测
        
            # 对检测后图片的处理
            r_image = detected[0]
            mkdir(predictFilePath +'output/')
            
            outputPath = predictFilePath +'output/' + str(i + 1) + '.jpg'
            
            r_image.save(outputPath)                         #保存检测过后的图片

            # 检测车流量
            new_boxs = detected[1]                           #新的boxes
            new_classes = detected[2]                        #新的classes
            traffic(old_boxs, new_boxs, old_classes, new_classes, image, i)
            old_boxs = new_boxs                              #变成老的boxes
            old_classes = new_classes                        #变成老的classes
            #
            fuckframejson[i]["count"][0]=traffic_count[2]#车辆
            fuckframejson[i]["count"][1]=traffic_count[0]#行人
            fuckframejson[i]["count"][2]=traffic_count[3]#摩托车
            fuckframejson[i]["count"][3]=traffic_count[7]#卡车

            #imagetest = Image.open(outputPath) #显示主图像-测试
            
            #imagetest.show()

            '''
            fuckjson = []                                        #存储所有物体的json文件的数组
            fuckjsoncar = []                                     #存储识别出车牌的车辆的json文件数组
            fuckframejson = []                                   #存储每帧所含物体的json文件数组
            此时这一帧的所有信息都已更新到数组内部了

            '''
            print(outputPath)                                         #当前帧检测结果图片
            print(fuckjson[i])                                        #当前帧所有物体的数据
            print(fuckjsoncar)                                     #当前帧所有车牌数据
            print(fuckframejson[i])                                   #当前帧所有车流量数据，以及所有


            
        #将json数组写入文件
        fr1 = open(jsonadd, 'w',encoding='utf-8') #打开json文件
        jsonstr = json.dumps(fuckjson)
        fr1.write(jsonstr)
        fr1.close()

        fr2 = open(carplatej, 'w',encoding='utf-8') #打开json文件
        jsonstr = json.dumps(fuckjsoncar)
        fr2.write(jsonstr)
        fr2.close()

        fr3 = open(jsonframe, 'w',encoding='utf-8') #打开json文件
        jsonstr = json.dumps(fuckframejson)
        fr3.write(jsonstr)
        fr3.close()
        for i in range(len(class_names)):          #遍历所有种类的数目，如果不等于零，则输出
            if traffic_count[i] != 0:
                print(class_names[i] + ': ' + str(traffic_count[i]))

    yolo.close_session()
    return True


